import os
from tree_sitter_c_config import rust_parser
from extract_component import extract_call_func_names
from global_config import CRATE_NAME
import logging

def split_rust_functions_in_file(file_path):
    if os.path.isfile(file_path):
        parser = rust_parser
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                code = f.read()
            # print(code)
            # 解析代码
            tree = parser.parse(bytes(code, 'utf8'))
            
            # 存储当前文件的函数
            functions = []

            # 使用深度优先搜索查找所有函数定义
            def traverse(node):
                if node.type == 'function_item':
                    start_byte = node.start_byte
                    end_byte = node.end_byte
                    function_code = code[start_byte:end_byte]
                    
                    # 获取函数名
                    for child in node.children:
                        if child.type == 'identifier':
                            function_name = code[child.start_byte:child.end_byte]
                            functions.append({
                                'name': function_name,
                                'code': function_code,
                                'start': start_byte,
                                'end': end_byte
                            })
                            break
                            
                for child in node.children:
                    traverse(child)
            
            traverse(tree.root_node)
            return functions
            
        except Exception as e:
            print(f"处理文件 {file_path} 时发生错误: {str(e)}")
    else:
        print(f"{file_path} 不是文件路径")
        return None


def get_dependency_ordered_functions(file_path):
    functions = split_rust_functions_in_file(file_path)
    base_result = functions
    if not functions:
        return []
    
    other_func_names = [f['name'] for f in functions if "test_" not in f['name']]
    sorted_functions = []

    for func in functions:
        res = extract_call_func_names(func['code'])
        func['dependency'] = set()
        for r in res:
            if r in other_func_names and r not in func['dependency'] and r != func['name']:
                func['dependency'].add(r)
        # print(f"{func['name']} : {func['dependency']}")

    other_funcs = [f for f in functions if "test_" not in f['name']]

    for func in other_funcs:
        name = func['name']
        for f in functions:
            if name in f['dependency']:
                f['dependency'] = f['dependency'] | func['dependency']

    # print("-" * 40)
    # for func in functions:
    #     print(f"{func['name']} : {func['dependency']}")
    test_funcs = [f for f in functions if "test_" in f['name']]
    
    while test_funcs:
        sorted_data = sorted(test_funcs, key=lambda x: len(x['dependency']))
        # names = [s['name'] for s in sorted_data]
        # print(names)
        if len(sorted_data[0]['dependency']) > 0:
            dependency = sorted_data[0]['dependency'].copy()
            for d in dependency:
                for func in test_funcs:
                    if d in func['dependency']:
                        func['dependency'].remove(d)
                for func in functions:
                    if func['name'] == d:
                        sorted_functions.append(func)
                        break
        sorted_functions.append(sorted_data[0])
        test_funcs.remove(sorted_data[0])
    # names = [s['name'] for s in sorted_functions]
    # print(names)


    return sorted_functions




# def split_rust_functions(folder_path):
#     """
#     使用tree-sitter解析Rust文件并提取所有函数
#     """
#     # 存储结果的字典
#     result = {}
     
#     # 遍历文件夹中的所有Rust文件
#     for root, dirs, files in os.walk(folder_path):
#         for file in files:
#             if file.endswith('.rs'):
#                 file_path = os.path.join(root, file)
#                 res = split_rust_functions_in_file(file_path=file_path)
#                 if res:
#                     result[file_path] = res
#                 else:
#                     print(f"处理文件{file_path}异常返回")
#     return result



# 'name': function_name
# 'code': function_code
# 'start': start_byte
# 'end': end_byte
def print_functions(functions_dict):
    """
    打印提取的函数信息
    
    Args:
        functions_dict: 从split_rust_functions返回的字典
    """
    for file_path, functions in functions_dict.items():
        print(f"\n文件: {file_path}")
        print("=" * 80)
        
        for func in functions:
            print(f"\n函数名: {func['name']}")
            print("-" * 40)
            print(func['code'])
            print("-" * 40)



# def save_functions_to_files(functions_dict, output_dir):
#     """
#     将提取的函数保存到单独的文件中
    
#     Args:
#         functions_dict: 从split_rust_functions返回的字典
#         output_dir: 输出目录的路径
#     """
#     # 确保输出目录存在
#     os.makedirs(output_dir, exist_ok=True)
    
#     for file_path, functions in functions_dict.items():
#         base_name = os.path.basename(file_path)
#         file_name_without_ext = os.path.splitext(base_name)[0]
        
#         file_dir = os.path.join(output_dir, file_name_without_ext)
#         os.makedirs(file_dir, exist_ok=True)
        
#         for func in functions:
#             # 清理文件名
#             safe_function_name = "".join(c for c in func['name'] if c.isalnum() or c in '_-')
#             if not safe_function_name:  # 如果函数名清理后为空，直接跳过
#                 continue
            
#             output_file = os.path.join(file_dir, f"{safe_function_name}.rs")
            
#             with open(output_file, 'w', encoding='utf-8') as f:
#                 f.write(func['code'])

# 修改使用示例:
# result = split_rust_functions(folder_path="/mnt/sda/xc/C2Rust/c_rust_agents/vivo-c2rust/cal-safe-ratio/init/")
# print_functions(result)  # 如果还需要打印的话
# save_functions_to_files(result, "Rust_functions")

if __name__ == "__main__":
    # result = split_rust_functions_in_file("/mnt/sda/xc/C2Rust/c_rust_agents/vivo-c2rust/Output/primary/src/avl_tree.rs")
    # for func in result:
    #     print(f"\n函数名: {func['name']}")
    #     print("-" * 40)
    #     print(func['code'])
    #     print("-" * 40)

    file_path = f"Output/{CRATE_NAME}/tests/test_avl_tree.rs"
    res = get_dependency_ordered_functions(file_path)
    for func in res:
        print(func['name'])